/*! @file DrawList.h
 *  @brief Data structure to store robot model to be drawn.
 *
 *  Stores all the data (except for joint positions) for the robot.
 *  Knows how to load cheetah robots from file.
 *  Will need to support adding a variable number of items (for loading terrain
 * which needs to be drawn)
 *
 *  Also supports having duplicated objects
 */

#ifndef PROJECT_DRAWLIST_H
#define PROJECT_DRAWLIST_H

#include "Checkerboard.h"
#include "Collision/CollisionPlane.h"
#include "Colors.h"
#include "Dynamics/DynamicsSimulator.h"
#include "Dynamics/FloatingBaseModel.h"
#include "Dynamics/spatial.h"
#include "SimUtilities/VisualizationData.h"
#include "cppTypes.h"
#include "obj_loader.h"
#include "sim_utilities.h"

#include <QMatrix4x4>

#include <stdlib.h>
#include <vector>
#include "Configuration.h"

class BoxInfo {
public:
    double depth, width, height;
    float frame[16];  // SE3
    /*
                       T[0] T[4] T[8]  T[12]
                       T[1] T[5] T[9]  T[13]
                       T[2] T[6] T[10] T[14]
                       T[3] T[7] T[11] T[15] = (0, 0, 0, 1)
                       */
};

struct ScrollInfo {
    size_t id;
    float xs, ys;
};

struct SolidColor {
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW
    Vec4<float> rgba;
    bool useSolidColor;
};

class DrawList {
public:
    VisualizationData *_visualizationData;

    DrawList() {
        _cameraOrigin = Vec3<double>::Zero();
        loadFiles();
    }
    size_t addCheetah3(Vec4<float> color, bool useOld, bool canHide);
    size_t addMiniCheetah(Vec4<float> color, bool useOld, bool canHide);
    void buildDrawList();
    void loadFiles();
    void doScrolling(Vec3<float> cameraPos);
    size_t addCheckerboard(Checkerboard &checkerBoard, bool setScrolling);
    size_t addDebugSphere(float radius);
    void addBox(double depth, double width, double height,
                const Vec3<double> &pos, const Mat3<double> &ori,
                bool transparent);
    void addMesh(double grid_size, const Vec3<double> &left_corner,
                 const DMat<double> &height_map, bool transparent);

    /*!
   * Resize to hold size objects.
   */
    void resize(size_t nUniqueObject, size_t nTotalObjects) {
        _nUnique = nUniqueObject;
        _nTotal = nTotalObjects;
        _vertexData.resize(nUniqueObject);
        _normalData.resize(nUniqueObject);
        _colorData.resize(nUniqueObject);
        _offsetXforms.resize(nUniqueObject);
        _objectMap.resize(nTotalObjects);
    }

    /*!
   * Get the total number of objects to be drawn
   */
    size_t getNumObjectsToDraw() { return _nTotal; }

    /*!
   * For the i-th object, get the offset into the model data.
   * For use with the glDrawArrays function.
   * Note that several objects may have the same geometry, so this will return
   * the same value for these objects!
   */
    size_t getGLDrawArrayOffset(size_t i) {
        return _glArrayOffsets.at(_objectMap.at(i));
    }

    /*!
   * For the i-th object, get the size of the model data array
   */
    size_t getGLDrawArraySize(size_t i) {
        return _glArraySizes.at(_objectMap.at(i));
    }

    bool getCanBeHidden(size_t i) {
        return _canBeHidden.at(i);
    }

    /*!
   * Get the array containing all vertex data.  Use getGLDrawArrayOffset/Size to
   * get indices and sizes for each object
   */
    float *getVertexArray() { return _glVertexData.data(); }

    /*!
   * Get the array containing all normal data.  Use getGLDrawArrayOffset/Size to
   * get indices and sizes for each object
   */
    float *getNormalArray() { return _glNormalData.data(); }

    size_t getSizeOfAllData() { return _glVertexData.size(); }

    /*!
   * Get the array containing all color data.
   */
    float *getColorArray() { return _glColorData.data(); }

    /*!
   * Get the Qt transformation matrix which should be applied to the model
   * geometry This is to correct for errors when exporting parts and shifts them
   * all to the origin.
   */
    QMatrix4x4 &getModelBaseTransform(size_t i) { return _modelOffsets[i]; }

    /*!
   * Get the Qt transformation matrix which should be applied to move a model
   * from the origin to where it should be in the world
   */
    QMatrix4x4 &getModelKinematicTransform(size_t i) {
        return _kinematicXform[i];
    }

    /*!
   * Get size of data used by the GPU in megabytes
   * For debugging
   */
    float getGLDataSizeMB() {
        size_t bytes =
                _glColorData.size() + _glNormalData.size() + _glVertexData.size();
        bytes = bytes * sizeof(float);
        return (float)bytes / float(1 << 20);
    }

    /*!
   * Returns true a single time if we have changed geometries and need to
   * reload.
   */
    bool needsReload() {
        if (_reloadNeeded) {
            _reloadNeeded = false;
            return true;
        }
        return false;
    }

    /*!
   * Update the position of a robot's bodies using the result of a dynamics
   * simulation. Doesn't run the simulator - just pulls _Xa from the
   * DynamicsSimulator
   * @param model  : the simulator
   * @param id     : the id returned from the loadCheetah3 or loadMiniCheetah
   * function.
   */
    template <typename T>
    void updateRobotFromModel(DynamicsSimulator<T> &model, size_t id,
                              bool updateOrigin = false) {
        for (size_t modelID = 5, graphicsID = id; modelID < model.getNumBodies();
             modelID++, graphicsID++) {
            _kinematicXform.at(graphicsID) =
                    spatialTransformToQT(model.getModel()._Xa.at(modelID));
        }

        if (updateOrigin) {
            _cameraOrigin = model.getState().bodyPosition.template cast<T>();
        }
    }

    /*!
   * Update the additional information drawn by GUI
   * Doesn't run the simulator
   *  - just pulls contact (or other in the future) data from the
   * DynamicsSimulator
   * @param model  : the simulator
   */
    template <typename T>
    void updateAdditionalInfo(DynamicsSimulator<T> &model)
    {
        if (_additionalInfoFirstVisit)
        {
            _nTotalGC = model.getTotalNumGC();
            _cp_touch.resize(_nTotalGC, false);
            _cp_pos.resize(_nTotalGC);
            _cp_force.resize(_nTotalGC);
            std::vector<double> tmp(3);
            for (size_t i(0); i < _nTotalGC; ++i)
            {
                _cp_pos[i] = tmp;
                _cp_force[i] = tmp;
            }
            _additionalInfoFirstVisit = false;
        }

        for (size_t i(0); i < _nTotalGC; ++i)
        {
            // TODO: check touch boolean
            _cp_touch[i] = true;
            for (size_t j(0); j < 3; ++j)
            {
                _cp_pos[i][j] = model.getModel()._pGC[i][j];
                _cp_force[i][j] = model.getContactForce(i)[j];
            }
        }
    }

    /*!
   * Updates the position of a checkerboard to match an infinite collision plane
   * The infinite collision plane only specifies orientation, so we
   * @param model : the collision plane
   * @param id    : the id retured from creating the checkerboard
   */
    template <typename T>
    void updateCheckerboardFromCollisionPlane(CollisionPlane<T> &model,
                                              size_t id)
    {
        // Mat4<T> H = sxformToHomogeneous(model.getLocation());
        _kinematicXform.at(id) = spatialTransformToQT(model.getLocation());
    }

    template <typename T>
    void updateCheckerboard(T height, size_t id)
    {
        QMatrix4x4 H;
        H.setToIdentity();
        H.translate(0., 0., height);
        _kinematicXform.at(id) = H;
    }

    template <typename T>
    void updateDebugSphereLocation(Vec3<T> &position, size_t id)
    {
        QMatrix4x4 H;
        H.setToIdentity();
        H.translate(position[0], position[1], position[2]);
        _kinematicXform.at(id) = H;
    }

    /*!
   * Fill color data with a solid color
   */
    static void setSolidColor(std::vector<float> &data,
                              size_t size,
                              float r,
                              float g,
                              float b)
    {
        data.clear();
        data.resize(size);

        if ((size % 3) != 0)
        {
            throw std::runtime_error("setSolidColor invalid size");
        }

        for (size_t i = 0; i < size / 3; i++)
        {
            data[i * 3 + 0] = r;
            data[i * 3 + 1] = g;
            data[i * 3 + 2] = b;
        }
    }

    /* Get Functions */
    const size_t &getTotalNumGC()
    {
        return _nTotalGC;
    }

    const std::vector<double> &getGCPos(size_t idx)
    {
        return _cp_pos[idx];
    }

    const std::vector<double> &getGCForce(size_t idx)
    {
        return _cp_force[idx];
    }

    const std::vector<BoxInfo> &getBoxInfoList()
    {
        return _box_list;
    }

    const DMat<double> &getHeightMap()
    {
        return _height_map;
    }

    const Vec3<double> &getHeightMapLeftCorner()
    {
        return _height_map_left_corner;
    }

    const double &getHeightMapMax()
    {
        return _height_map_max;
    }

    const double &getHeightMapMin()
    {
        return _height_map_min;
    }

    const double &getGridSize()
    {
        return _grid_size;
    }

    const Vec3<double> &getCameraOrigin()
    {
        return _cameraOrigin;
    }

    vectorAligned<SolidColor> _instanceColor;
    std::vector<QMatrix4x4> _kinematicXform;

private:
    size_t _nUnique = 0, _nTotal = 0;
    std::vector<std::vector<float>> _vertexData;
    std::vector<std::vector<float>> _normalData;
    std::vector<std::vector<float>> _colorData;
    std::vector<u8> _canBeHidden;
    vectorAligned<Mat4<float>> _offsetXforms;  // these are NOT coordinate transformations!
    std::string _baseFileName = THIS_COM"./resources/"; /* "../resources/"*/

    std::vector<size_t> _objectMap;

    std::vector<size_t> _glArrayOffsets;
    std::vector<size_t> _glArraySizes;

    std::vector<float> _glVertexData;
    std::vector<float> _glNormalData;
    std::vector<float> _glColorData;
    std::vector<ScrollInfo> _scrollIDs;

    std::vector<QMatrix4x4> _modelOffsets;

    bool _reloadNeeded = false;
    bool _additionalInfoFirstVisit = true;

    size_t _nTotalGC = 0;
    std::vector<bool> _cp_touch;
    std::vector<std::vector<double>> _cp_pos;
    std::vector<std::vector<double>> _cp_force;
    std::vector<BoxInfo> _box_list;

    double _grid_size;
    Vec3<double> _height_map_left_corner;
    DMat<double> _height_map;
    double _height_map_max, _height_map_min;

    Vec3<double> _cameraOrigin;

    size_t _cheetah3LoadIndex = 0, _miniCheetahLoadIndex = 0,
    _sphereLoadIndex = 0, _cubeLoadIndex = 0;
};

#endif  // PROJECT_DRAWLIST_H
